import traceback
import cv2
import numpy as np
import os
import glob
import zxingcpp
from urllib.parse import urlparse, parse_qs
from tqdm import tqdm
from PIL import Image
from PIL.ExifTags import TAGS
from pathlib import Path
from PIL import UnidentifiedImageError
from mimetypes import guess_type
import concurrent.futures

from pillow_heif import register_heif_opener

register_heif_opener()


def read_image(image_path, mode=cv2.IMREAD_COLOR):
    """
    读取图片，返回图片对象，支持中文路径
    """
    # 解决cv2.imread不支持中文路径的问题
    img = Image.open(image_path)
    # return np.array(image)
    # img_array = np.fromfile(image_path, dtype=np.uint8)
    # img = cv2.imdecode(img_array, mode)
    if img is None:
        raise ValueError("图片路径错误")
    return img


def detect_qrcodes_from_img(img):
    """
    检测图片对象中的二维码，返回二维码内容列表
    """
    img = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)
    barcodes = zxingcpp.read_barcodes(img)
    if barcodes:
        return [barcode.text for barcode in barcodes if barcode.text]
    return []


def contains_watermark_from_img(img, template, threshold=0.8):
    """
    检测图片对象中是否包含水印模板
    """
    return False

    # 使用多尺度模板匹配以适应水印大小变化
    found = False
    for scale in np.linspace(0.5, 1.5, 11):
        h, w = template.shape[:2]
        resized = cv2.resize(
            template, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA
        )
        if resized.shape[0] > img.shape[0] or resized.shape[1] > img.shape[1]:
            continue
        res = cv2.matchTemplate(img, resized, cv2.TM_CCOEFF_NORMED)
        loc = np.where(res >= threshold)
        if len(list(zip(*loc[::-1]))) > 0:
            found = True
            break
    return found


def detect_image(img):
    # 只取图片的下半部分进行二维码检测
    # h = img.shape[0]
    # img_lower_half = img[h // 2 :, :]
    qrcodes = detect_qrcodes_from_img(img)
    if qrcodes:
        # 整理二维码内容，若为URL则拆分为uri和参数map
        result = {"qrcodes": []}
        for code in qrcodes:
            if code.startswith("http://") or code.startswith("https://"):
                parsed = urlparse(code)
                uri = f"{parsed.scheme}://{parsed.netloc}{parsed.path}"
                params = {
                    k: v[0] if len(v) == 1 else v
                    for k, v in parse_qs(parsed.query).items()
                }
                result["qrcodes"].append({"uri": uri, "params": params})
            else:
                result["qrcodes"].append({"text": code})
        return result

    watermark_dir = "water_mark"
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    watermark = None

    global watermark_templates
    if "watermark_templates" not in globals():
        watermark_templates = {}
        for fname in os.listdir(watermark_dir):
            tpath = os.path.join(watermark_dir, fname)
            if not os.path.isfile(tpath):
                continue
            watermark_templates[fname] = read_image(tpath, cv2.IMREAD_GRAYSCALE)

    for filename, template in watermark_templates.items():
        # template = read_image(template_path, cv2.IMREAD_GRAYSCALE)
        if contains_watermark_from_img(img_gray, template):
            watermark = filename.split(".")[0]
            break

    return {"watermark": watermark}


def extract_main_site(url):
    """
    解析url，得到主站名字，比如www.bilibili.com就是bilibili
    """
    parsed = urlparse(url)
    hostname = parsed.hostname
    if not hostname:
        return None
    # 去除www.前缀
    if hostname.startswith("www."):
        hostname = hostname[4:]
    # 取第一个域名部分
    parts = hostname.split(".")
    if len(parts) >= 2:
        return parts[-2]
    return hostname


def detect_image_source(img):
    result = detect_image(img)
    if "watermark" in result and result["watermark"]:
        return result["watermark"]
    if "qrcodes" in result and result["qrcodes"]:
        for code in result["qrcodes"]:
            if "uri" in code:
                return extract_main_site(code["uri"])
    return None


def perform_image_test(func):
    img_extensions = (".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp")
    test_dir = "test_imgs"
    img_files = []
    for ext in img_extensions:
        img_files.extend(
            glob.glob(os.path.join(test_dir, "**", f"*{ext}"), recursive=True)
        )

    for img_path in img_files:
        try:
            img = read_image(img_path)
            result = func(img)
            print(f"{img_path}: {result}")
        except Exception as e:
            print(f"{img_path}: Error - {e}")


def judge_dest_dir(img_path):
    """
    判断图片来源，返回来源字符串
    """
    try:
        ext = os.path.splitext(img_path)[1].lower()
        print(f"ext: {ext}")
        if ext == ".mp4":
            return "mp4"
        elif ext == ".gif":
            return "gif"
        else:
            img = read_image(img_path)
            if is_camera_photo(img):
                return "camera"
            else:
                source = detect_image_source(img)
                if not source:
                    return "unknown"
                return source
    except Exception as e:
        print(f"Error processing {img_path}: {e}")
        print(traceback.format_exc())
        return None


def move_images_by_source_single(args):
    img_path, dst_base_dir = args
    if not os.path.isfile(img_path):
        return f"{img_path}: Not a file"
    try:

        source = judge_dest_dir(img_path)
        dst_dir = os.path.join(dst_base_dir, source)
        os.makedirs(dst_dir, exist_ok=True)
        dst_path = os.path.join(dst_dir, os.path.basename(img_path))
        if os.path.exists(dst_path):
            # 如果目标文件已存在，跳过
            # return f"{img_path}: Already exists at {dst_path}"
            print(f"{img_path}: Already exists at {dst_path}")
            return None
        # 保留时间戳
        stat = os.stat(img_path)
        os.rename(img_path, dst_path)
        os.utime(dst_path, (stat.st_atime, stat.st_mtime))
        # print(f"Moved {img_path} -> {dst_path}")
        return None
    except Exception as e:
        return f"{img_path}: Error - move_images_by_source_single - {e} \n{traceback.format_exc()}"


def is_image_file(path):
    try:
        with Image.open(path) as img:
            img.verify()
        return True
    except (UnidentifiedImageError, OSError):
        return False


def has_image_suffix(filename):
    mime, _ = guess_type(filename)
    return mime is not None and mime.startswith("image/")


def glob_source_files(src_dir):
    img_extensions = (
        ".jpg",
        ".jpeg",
        ".png",
        ".bmp",
        ".tiff",
        ".webp",
        ".heic",
        ".heif",
        ".mp4",
        ".gif",
    )
    encountered_suffix = set()
    img_files = []
    for each_file in Path(src_dir).rglob("*"):
        encountered_suffix.add(each_file.suffix.lower())
        if each_file.suffix.lower() in img_extensions and each_file.is_file():
            # if each_file.suffix.lower() in img_extensions:
            img_files.append(str(each_file))
    print(f"encountered_suffix: {encountered_suffix}")
    return img_files


def process_all_image(src_dir, map_func):
    dst_base_dir = src_dir
    os.rename(src_dir, src_dir + ".remain")
    src_dir += ".remain"

    img_files = glob_source_files(src_dir)
    args_list = [(img_path, dst_base_dir) for img_path in img_files]

    with concurrent.futures.ThreadPoolExecutor() as executor:
        results = list(
            tqdm(
                executor.map(map_func, args_list),
                total=len(img_files),
                desc="Moving images",
            )
        )
        for res in results:
            if res:
                print(res)
    # 递归清理空文件夹
    for root, dirs, files in os.walk(src_dir, topdown=False):
        for d in dirs:
            dir_path = os.path.join(root, d)
            try:
                if not os.listdir(dir_path):
                    os.rmdir(dir_path)
            except Exception as e:
                print(f"Failed to remove {dir_path}: {e}")


def is_camera_photo(img):
    """
    判断图片是否为摄像机拍照（通过EXIF数据分析），返回True/False/None
    """
    try:
        exif_data = img.getexif()
        if not exif_data:
            # print("没有EXIF信息")
            return None  # 没有EXIF信息，无法判断

        exif = {TAGS.get(k, k): v for k, v in exif_data.items()}
        # 摄像机拍照通常有Make/Model/DateTimeOriginal等字段
        # print(exif)
        camera_tags = ["Make", "Model", "DateTimeOriginal", "LensModel"]
        if any(tag in exif for tag in camera_tags):
            return True
        # 某些系统生成图片可能有"Software"字段，且无上述字段
        # if "Software" in exif and not any(tag in exif for tag in camera_tags):
        #     return False
        # return None
    except Exception as e:
        print(f"EXIF解析失败: {e}")
        return None


# 示例用法
if __name__ == "__main__":
    # perform_image_test(detect_image_source)
    process_all_image(
        r"F:\Picture 图片\来自手机\others",
        move_images_by_source_single,
    )
    # TODO mp4 归为一类 gif归为一类

    # read_image(
    #     r"F:\Picture 图片\来自手机\others.remain\有链接图片分类\camera\IMG_20250124_124836.HEIC"
    # )
    # # files = glob_img_files(
    #     r"F:\Picture 图片\来自手机\来自：Redmi K80.remain.remain.remain"
    # )
    # print(files)
    # judge_image_source(
    #     r"F:\Picture 图片\来自手机\来自：Redmi K80\相机\MVIMG_20250514_131554.jpg"
    # )
    # img = Image.open(
    #     r"F:\Picture 图片\来自手机\来自：Redmi K80\相机\MVIMG_20250514_131554.jpg"
    # )
    # is_camera_photo(img)
    # 示例：检测test_imgs目录下图片是否为摄像机拍照
    # test_dir = "test_imgs"
    # img_extensions = (".jpg", ".jpeg", ".png")
    # img_files = []
    # for ext in img_extensions:
    #     img_files.extend(glob.glob(os.path.join(test_dir, "**", f"*{ext}"), recursive=True))

    # for img_path in img_files:
    #     result = is_camera_photo(img_path)
    #     print(f"{img_path}: {'Camera Photo' if result else 'System Generated' if result is False else 'Unknown'}")
